package com.zhang.sparksql_1

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.{Aggregator, MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, LongType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Encoder, Encoders, Row, SparkSession, functions}

/**
 * @title:
 * @author: zhang
 * @date: 2021/12/10 19:48 
 */
object SparkSQl_UDAF_04 {
  def main(args: Array[String]): Unit = {
    //  获取SparkSession连接对象
    val conf = new SparkConf().setMaster("local[*]").setAppName("spark-sql")
    val spark: SparkSession = SparkSession.builder().config(conf).getOrCreate()
    //todo 创建DataFrame
    val df: DataFrame = spark.read.json("datas/user.json")

    df.createOrReplaceTempView("user")

    //todo 创建UDAF函数 强类型
    spark.udf.register("MyAvg",functions.udaf(new MyAvgUDAF()))
    spark.sql("select MyAvg(age) from user").show()
    //todo 关闭资源
    spark.stop()
  }
//自定义聚合函数类：计算年龄的平均值
  // IN ：输入的数据类型Long
  // BUF：缓冲区的数据类型Buff
  //OUT：输出的数据类型Long

  case class Buff(var total:Long,var count:Long)
  class MyAvgUDAF extends Aggregator[Long,Buff,Long] {
    override def zero: Buff = {
      Buff(0L,0L)
    }

    override def reduce(buff: Buff, in: Long): Buff = {
      buff.total = buff.total+in
      buff.count+=1
      buff
    }

    override def merge(buff1: Buff, buff2: Buff): Buff = {
      buff1.total+=buff2.total
      buff1.count+=buff2.count
      buff1
    }

    override def finish(buff: Buff): Long = {
      buff.total / buff.count
    }

    override def bufferEncoder: Encoder[Buff] = Encoders.product

    override def outputEncoder: Encoder[Long] = Encoders.scalaLong
  }


}
